"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2025/3/6 16:59
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    utils.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    utils_2025/3/6.0.1

-------------------------------------------------------------------------------
# @Description:

    ---------------------------------------------------------------------------
    [Update History]:
        2025/3/6: impl FisheyeCrop class and timer

-------------------------------------------------------------------------------
"""
import cv2
import numpy as np
import time
import statistics
from typing import Dict, List, Optional
from functools import wraps


class TimingAnalyzer:
    """
    函数耗时分析器，用于精确测量函数执行时间并进行统计分析
    """

    def __init__(self):
        self.timing_data: Dict[str, List[float]] = {}
        self.call_counts: Dict[str, int] = {}
        self.is_active = True

    def start_timer(self, name: str) -> float:
        """开始计时特定名称的操作"""
        if not self.is_active:
            return 0.0

        if name not in self.timing_data:
            self.timing_data[name] = []
            self.call_counts[name] = 0

        self.call_counts[name] += 1
        return time.perf_counter()

    def stop_timer(self, name: str, start_time: float) -> float:
        """停止计时并记录耗时"""
        if not self.is_active or start_time == 0.0:
            return 0.0

        elapsed = time.perf_counter() - start_time
        self.timing_data[name].append(elapsed)
        return elapsed

    def get_stats(self, name: str) -> Dict[str, float]:
        """获取特定操作的统计数据"""
        if name not in self.timing_data or not self.timing_data[name]:
            return {
                "calls": 0,
                "total": 0.0,
                "avg": 0.0,
                "min": 0.0,
                "max": 0.0,
                "median": 0.0,
                "std_dev": 0.0
            }

        times = self.timing_data[name]
        return {
            "calls": self.call_counts[name],
            "total": sum(times),
            "avg": statistics.mean(times),
            "min": min(times),
            "max": max(times),
            "median": statistics.median(times),
            "std_dev": statistics.stdev(times) if len(times) > 1 else 0.0
        }

    def get_all_stats(self) -> Dict[str, Dict[str, float]]:
        """获取所有操作的统计数据"""
        return {name: self.get_stats(name) for name in self.timing_data}

    def reset(self) -> None:
        """重置所有计时数据"""
        self.timing_data.clear()
        self.call_counts.clear()

    def print_report(self, sort_by: str = "total") -> None:
        """打印性能报告"""
        all_stats = self.get_all_stats()
        if not all_stats:
            print("没有计时数据")
            return

        # 按指定字段排序
        sorted_stats = sorted(
            all_stats.items(),
            key=lambda x: x[1][sort_by],
            reverse=True
        )

        print(f"\n{'=' * 140}")
        print(f"{'Func name':50} | {'Call times':12} | {'Total tims(s)':14} | {'Mean time(s)':14} | {'Minium(s)':14} | {'Maximum(s)':14}")
        print(f"{'-' * 140}")

        for name, stats in sorted_stats:
            print(f"{name:50} | {stats['calls']:12d} | {stats['total']:14.6f} | "
                  f"{stats['avg']:14.6f} | {stats['min']:14.6f} | {stats['max']:14.6f}")

        print(f"{'=' * 140}\n")

    def time_function(self, func=None, name: Optional[str] = None):
        """函数装饰器，用于自动计时函数执行时间"""

        def decorator(fn):
            @wraps(fn)
            def wrapper(*args, **kwargs):
                fn_name = name or fn.__qualname__
                start = self.start_timer(fn_name)
                try:
                    result = fn(*args, **kwargs)
                    return result
                finally:
                    elapsed = self.stop_timer(fn_name, start)
                    # print(f" [Info] {fn_name} Elapsed time: {elapsed * 1000:.6f} ms")

            return wrapper

        if func is None:
            return decorator
        return decorator(func)

    def time_block(self, name: str):
        """上下文管理器，用于计时代码块执行时间"""

        class TimingContext:
            def __init__(self, analyzer, block_name):
                self.analyzer = analyzer
                self.name = block_name
                self.start_time = 0.0

            def __enter__(self):
                self.start_time = self.analyzer.start_timer(self.name)
                return self

            def __exit__(self, exc_type, exc_val, exc_tb):
                self.analyzer.stop_timer(self.name, self.start_time)

        return TimingContext(self, name)


# Initial timer
timer = TimingAnalyzer()


class FisheyeCrop:
    def __init__(self, bin_th: float = 5, mor_ksize: int = 5, valid_len: int = 4):
        self.bin_th = bin_th
        self.mor_ksize = mor_ksize
        self.valid_len = valid_len
        self.target_crop_param = None

    @timer.time_function
    def find_crop_index(self, binary: np.ndarray):
        height, width = binary.shape
        c_y, c_x = height // 2, width // 2

        d_values = []
        for col_offset in range(-4, 4, 2):
            col = c_x + col_offset

            y_u = None
            for row in range(0, c_y):
                if binary[row, col] > 0:
                    consecutive_count = 0
                    for check_y in range(row, min(row + 5, c_y)):
                        if binary[check_y, col] > 0:
                            consecutive_count += 1
                    if consecutive_count >= self.valid_len:
                        y_u = row
                        break

            y_d = None
            for row in range(height - 1, c_y, -1):
                if binary[row, col] > 0:
                    consecutive_count = 0
                    for check_y in range(max(row - 5, c_y), row):
                        if binary[check_y, col] > 0:
                            consecutive_count += 1
                    if consecutive_count >= self.valid_len:
                        y_d = row
                        break
            if y_u is not None and y_d is not None:
                d_values.append(y_d - y_u)

        for row_offset in range(-4, 4, 2):
            row = c_y + row_offset

            x_l = None
            for col in range(0, c_x):
                if binary[row, col] > 0:
                    consecutive_count = 0
                    for check_x in range(col, min(col + 5, c_x)):
                        if binary[row, check_x] > 0:
                            consecutive_count += 1
                    if consecutive_count >= self.valid_len:
                        x_l = col
                        break

            x_r = None
            for col in range(width - 1, c_x, -1):
                if binary[row, col] > 0:
                    consecutive_count = 0
                    for check_x in range(max(col - 5, c_x), col):
                        if binary[row, check_x] > 0:
                            consecutive_count += 1
                    if consecutive_count >= self.valid_len:
                        x_r = col
                        break
            if x_l is not None and x_r is not None:
                d_values.append(x_r - x_l)

        if not d_values:
            return 0, width, 0, height

        d_values.sort()
        diameter = d_values[len(d_values) * 3 // 4]  # 取 75% 分位数

        # 计算裁剪区域
        x_l = max(0, c_x - diameter // 2)
        x_r = min(width, c_x + diameter // 2)
        y_u = max(0, c_y - diameter // 2)
        y_d = min(height, c_y + diameter // 2)

        # 确保裁剪区域是方形
        crop_size = min(x_r - x_l, y_d - y_u)

        x_l = c_x - crop_size // 2
        x_r = c_x + crop_size // 2
        y_u = c_y - crop_size // 2
        y_d = c_y + crop_size // 2

        # 确保索引在有效范围内
        x_l = max(0, x_l)
        x_r = min(width, x_r)
        y_u = max(0, y_u)
        y_d = min(height, y_d)

        return x_l, x_r, y_u, y_d

    @timer.time_function
    def find_crop_index_fast(self, binary: np.ndarray):
        height, width = binary.shape
        c_y, c_x = height // 2, width // 2

        d_values = []
        # 保持原有的采样范围，但使用numpy操作优化
        for col_offset in range(-4, 4, 2):
            col = c_x + col_offset

            # 垂直方向扫描优化
            col_data_up = binary[:c_y, col]
            col_data_down = binary[c_y:, col]

            y_u = None
            for i in range(len(col_data_up)):
                if col_data_up[i] > 0:
                    window = col_data_up[i:min(i + 5, c_y)]
                    if np.sum(window > 0) >= self.valid_len:
                        y_u = i
                        break

            y_d = None
            for i in range(len(col_data_down) - 1, -1, -1):
                if col_data_down[i] > 0:
                    window = col_data_down[max(0, i - 4):i + 1]
                    if np.sum(window > 0) >= self.valid_len:
                        y_d = c_y + i
                        break

            if y_u is not None and y_d is not None:
                d_values.append(y_d - y_u)

        # 水平方向扫描保持相同的逻辑优化
        for row_offset in range(-4, 4, 2):
            row = c_y + row_offset

            row_data_left = binary[row, :c_x]
            row_data_right = binary[row, c_x:]

            x_l = None
            for i in range(len(row_data_left)):
                if row_data_left[i] > 0:
                    window = row_data_left[i:min(i + 5, c_x)]
                    if np.sum(window > 0) >= self.valid_len:
                        x_l = i
                        break

            x_r = None
            for i in range(len(row_data_right) - 1, -1, -1):
                if row_data_right[i] > 0:
                    window = row_data_right[max(0, i - 4):i + 1]
                    if np.sum(window > 0) >= self.valid_len:
                        x_r = c_x + i
                        break

            if x_l is not None and x_r is not None:
                d_values.append(x_r - x_l)

        if not d_values:
            return 0, width, 0, height

        # 保持原有的75%分位数选择
        d_values.sort()
        diameter = d_values[len(d_values) * 3 // 4]

        # 计算裁剪区域
        x_l = max(0, c_x - diameter // 2)
        x_r = min(width, c_x + diameter // 2)
        y_u = max(0, c_y - diameter // 2)
        y_d = min(height, c_y + diameter // 2)

        # 确保裁剪区域是方形
        crop_size = min(x_r - x_l, y_d - y_u)

        x_l = c_x - crop_size // 2
        x_r = c_x + crop_size // 2
        y_u = c_y - crop_size // 2
        y_d = c_y + crop_size // 2

        # 确保索引在有效范围内
        x_l = max(0, int(x_l))
        x_r = min(width, int(x_r))
        y_u = max(0, int(y_u))
        y_d = min(height, int(y_d))

        return x_l, x_r, y_u, y_d

    @timer.time_function
    def crop(self, img: np.ndarray, is_first: bool = True, en_mor: bool = False) -> np.ndarray:
        if not is_first and self.target_crop_param is not None:
            x_l, x_r, y_u, y_d = self.target_crop_param
            cropped_img = img[y_u:y_d, x_l:x_r]
            return cropped_img

        if len(img.shape) == 3:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        else:
            gray = img

        _, binary = cv2.threshold(gray, self.bin_th, 255, cv2.THRESH_BINARY)

        if en_mor:
            kernel = np.ones((self.mor_ksize, self.mor_ksize), np.uint8)
            binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

        # find radius
        crop_index = self.find_crop_index(binary)
        # crop_index_1 = self.find_crop_index_fast(binary)
        if crop_index[0] == 0 and crop_index[2] == 0:
            self.target_crop_param = None
        else:
            self.target_crop_param = crop_index

        x_l, x_r, y_u, y_d = crop_index
        cropped_img = img[y_u:y_d, x_l:x_r]
        return cropped_img


class FisheyeUndistorter:
    """
        Corrects fisheye distortion using a geometric projection model.
    Args:
        fov_degrees:   Field of view of the fisheye camera in degrees.
        fisheye_img_w: Width of the original fisheye image.
        fisheye_img_h: Height of the original fisheye image.
        undistorted_w: Width of the undistorted output image.
    """

    def __init__(
            self,
            fov_degrees: float,
            fisheye_img_w: int,
            fisheye_img_h: int,
            undistorted_w: int
    ):
        if not (0 < fov_degrees < 360):
            raise ValueError("fov_degrees must be in (0, 180)")

        if fisheye_img_w <= 0 or fisheye_img_h <= 0 or undistorted_w <= 0:
            raise ValueError("Image dimensions must be positive integers")

        if fisheye_img_w != fisheye_img_h:
            raise ValueError("Image dimensions must be equal to image dimensions")

        self.fisheye_w = fisheye_img_w
        self.fisheye_h = fisheye_img_h
        self.undistorted_w = undistorted_w
        self.fov_rad = np.deg2rad(fov_degrees)

        # 生成映射关系
        self.target_x_coords, self.target_y_coords, self.radius = self.create_fisheye_projection_map(
            undistorted_w=self.undistorted_w,
            undistorted_h=self.fisheye_h,
            target_w=self.fisheye_w,
            target_h=self.fisheye_h
        )

    def create_fisheye_projection_map(
            self,
            undistorted_w: int,  # 去扭曲图像（源）的宽度 (Ws)
            undistorted_h: int,  # 去扭曲图像（源）的高度 (Hs)
            target_w: int,  # 鱼眼图像（目标）的宽度 (Wd)
            target_h: int,  # 鱼眼图像（目标）的高度 (Hd)
    ):
        """
            Generate coordinate maps for projecting a rectangular image to fisheye.
        """
        # 生成源图像的像素网格坐标（原点在左上角）
        y_pixels, x_pixels = np.indices((undistorted_h, undistorted_w), dtype=np.float32)

        # 将坐标原点移至图像中心，并调整方向（y轴向上为正）
        y_centered = undistorted_h / 2.0 - y_pixels  # 转换为笛卡尔坐标系
        x_centered = x_pixels - undistorted_w / 2.0

        # 将平面坐标转换为球面坐标（等距投影模型）
        # theta: 方位角（绕z轴的旋转角，范围 [-π, π]）
        # phi: 极角（与z轴的夹角，范围 [0, fov/2]）
        theta, phi = self.cartesian_to_spherical(x_centered, y_centered, undistorted_w, undistorted_h)

        # 计算鱼眼图像中的极坐标半径（等距投影：半径与极角phi成正比）
        radius_from_center = target_h * phi / self.fov_rad

        # 将极坐标转换为鱼眼图像的笛卡尔坐标（原点在图像中心）
        x_fisheye = radius_from_center * np.cos(theta)
        y_fisheye = radius_from_center * np.sin(theta)

        # 调整坐标系到目标图像的左上角原点（OpenCV标准坐标系）
        target_y_coords = target_h / 2.0 - y_fisheye
        target_x_coords = target_w / 2.0 + x_fisheye

        return target_x_coords, target_y_coords, radius_from_center

    def cartesian_to_spherical(self, x_centered, y_centered, undistorted_w, undistorted_h):
        """
            Return the equirectangular projection on a unit sphere,
            given cartesian coordinates of the de-warped image.
        """
        theta_alt = x_centered * self.fov_rad / undistorted_w
        phi_alt = y_centered * np.pi / undistorted_h

        x = np.sin(theta_alt) * np.cos(phi_alt)
        y = np.sin(phi_alt)
        z = np.cos(theta_alt) * np.cos(phi_alt)

        return np.arctan2(y, x), np.arctan2(np.sqrt(x ** 2 + y ** 2), z)

    @timer.time_function
    def apply_undistortion(self, img: np.ndarray) -> np.ndarray:
        """执行鱼眼校正"""
        return cv2.remap(img, self.target_x_coords, self.target_y_coords, cv2.INTER_LINEAR)


if __name__ == '__main__':
    """ [Config param] """
    fov_degree = 195
    img_left_path = r"C:\ws_rk\ws_2025\image-stitching\test\a.png"
    img_right_path = r"C:\ws_rk\ws_2025\image-stitching\test\b.png"

    img_ori_l = cv2.imread(img_left_path)
    img_ori_r = cv2.imread(img_right_path)

    # Crop fisheye image
    cropper = FisheyeCrop()
    fe_img_l = cropper.crop(img_ori_l, is_first=True, en_mor=False)
    fe_img_r = cropper.crop(img_ori_r, is_first=False, en_mor=False)

    if fe_img_l.shape != fe_img_r.shape:
        raise ValueError("Left image dimensions must be equal to right image dimensions!")

    image_w = fe_img_l.shape[1]
    image_h = fe_img_l.shape[0]
    undistort_image_w = int(fov_degree / 180 * image_h)
    stitched_image_w = image_w * 2

    # Undistort image
    undistorter = FisheyeUndistorter(fov_degree, image_w, image_h, undistort_image_w)
    undistort_img_l = undistorter.apply_undistortion(fe_img_l)
    undistort_img_r = undistorter.apply_undistortion(fe_img_r)

    timer.print_report()
    print('Done')
